from langchain_zhipu import ChatZhipuAI
from WriteDocuments import FAISSService
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain
from langchain.chains.retrieval import create_retrieval_chain
class ChatbotWithRetrieval:
    def __init__(self, base_dir):
        model = ChatZhipuAI(api_key="9b1f2582ae0c44c036ba1fdabed75c7b.FsfpGGSXKU4EZdsf"
                            , model="glm-4", verbose=True)
        print("执行加载数据")
        model.do_sample = False  # 温度设置为0，结果随机性 ghbnm
        faiss = FAISSService()
        print("执行加载完数据")

        vectorstore = faiss.vectorstore

        prompt = ChatPromptTemplate.from_template("""
        System: You are assistant ,chinese name is 吉利AI助手，We need you to respond to various questions from users
        System: The context includes some customer Q&A records and product information,
        Answer the following questions based on context and user's historical questions and assistant's answers,
        System: When it comes to time, please verify the content in the context based on the current time，
        System: current time {current_time}
        Please answer in Chinese:

        <context>
        {context}
        </context>
        
        <history>
        {history}
        </history>
        Question: {input}
        
        System: Priority should be given to using the given information for answering. 
        If the given information is unreasonable or you do not know the answer, please use your own knowledge base to answer. If you don't know, please answer truthfully and don't guess.
        System: If possible, please use a universal knowledge base to verify the given information and provide accurate and reasonable answers based on universal knowledge.
        System: Please answer as concisely as possible""")
        document_chain = create_stuff_documents_chain(model, prompt)

        retriever = vectorstore.as_retriever()
        retrieval_chain = create_retrieval_chain(retriever, document_chain)

        self.store = faiss
        # 向量数据库
        self.vectorstore = vectorstore

        # 初始化LLM
        self.llm = model

        # 初始化Memory
        self.memory = ""
        # 初始化对话历史
        self.conversation_history = ""
        self.qa = retrieval_chain
